import dataclasses
from typing import List
import marshal
import torch
@dataclasses.dataclass
class TestRequest:
    """
    TestRequest: A request for testing the server's performance
    """
    
    prompt: str
    prompt_len: int
    output_len: int
    
@dataclasses.dataclass
class TraceRequest:
    prompt:str
    prompt_len:int
    output_len:int
    interval:float
@dataclasses.dataclass
class Dataset:
    """
    Dataset: A dataset for testing the server's performance
    """
 
    dataset_name: str	# "sharegpt" / "alpaca" / ...
    reqs: List[TestRequest|TraceRequest]
    
    def dump(self, output_path: str):
        if isinstance(self.reqs[0],TestRequest):
            marshal.dump({
                "dataset_name": self.dataset_name,
                "reqs": [(req.prompt, req.prompt_len, req.output_len) for req in self.reqs]
            }, open(output_path, "wb"))
        else:
            marshal.dump({
                "dataset_name": self.dataset_name,
                "reqs": [(req.prompt, req.prompt_len, req.output_len,req.interval) for req in self.reqs]
            }, open(output_path, "wb"))
    
    @staticmethod
    def load(input_path: str):
        loaded_data = marshal.load(open(input_path, "rb"))
        def build_request(req):
            if len(req) == 3:
                return TestRequest(req[0], req[1], req[2])
            elif len(req) == 4:
                return TraceRequest(req[0], req[1], req[2], req[3])
            else:
                raise ValueError(f"Invalid request format: {req}")
        
        return Dataset(
            loaded_data["dataset_name"],
            [build_request(req) for req in loaded_data["reqs"]]
        )
        
import dataclasses
import numpy as np
from typing import List
import json

class RequestResult:
    """
    A class for storing the results of a single request
    """
    
    def __init__(
        self,
        prompt_len: int,
        output_len: int,
        start_time: float,
        end_time: float,
    ):
        self.prompt_len = prompt_len
        self.output_len = output_len
        self.start_time = start_time
        self.end_time = end_time

        
        self.latency = end_time - start_time



def read_request_results(path: str) -> List[RequestResult]:
    with open(path, "r") as f:
        request_results: List[RequestResult] = [
            RequestResult(
                item["prompt_len"],
                item["output_len"],
                item["start_time"],
                item["end_time"]
            )
            for item in json.load(f)
        ]
    return request_results

def count_valid_results(request_results: list[RequestResult], ftl: float, tpot: float) -> int:
    """
    count_valid_results: Count the number of requests that satisfy the given FTL and TPOT.
    """
    count = 0
    for req in request_results:
        if req.ftl <= ftl and req.tpot <= tpot:
            count += 1
    return count

def get_slo_attainment(request_results: list[RequestResult], ftl: float, tpot: float) -> float:
    """
    get_slo_attainment: Get the SLO attainment of the given request results under the given FTL and TPOT.
    """
    return count_valid_results(request_results, ftl, tpot) / len(request_results)
